import heapq
import bitarray

class Node:
    def __init__(self, freq, symbol=None, left=None, right=None):
        self.freq = freq
        self.symbol = symbol
        self.left = left
        self.right = right

    def __lt__(self, other):
        return self.freq < other.freq

def count_frequencies(data):
    frequencies = {}
    for symbol in data:
        if symbol in frequencies:
            frequencies[symbol] += 1
        else:
            frequencies[symbol] = 1
    return frequencies

def build_huffman_tree(frequencies):
    heap = [Node(freq, sym) for sym, freq in frequencies.items()]
    heapq.heapify(heap)

    while len(heap) > 1:
        node1 = heapq.heappop(heap)
        node2 = heapq.heappop(heap)
        merged_node = Node(node1.freq + node2.freq, left=node1, right=node2)
        heapq.heappush(heap, merged_node)

    return heap[0]

def build_huffman_table(node, code='', huffman_table={}):
    if node.left is None and node.right is None:
        huffman_table[node.symbol] = code
    if node.left:
        build_huffman_table(node.left, code + '0', huffman_table)
    if node.right:
        build_huffman_table(node.right, code + '1', huffman_table)

def compress(input_file, output_file):
    with open(input_file, 'r') as file:
        data = file.read()

    frequencies = count_frequencies(data)
    huffman_tree = build_huffman_tree(frequencies)
    huffman_table = {}
    build_huffman_table(huffman_tree, huffman_table=huffman_table)

    encoded_data = bitarray.bitarray()
    for symbol in data:
        encoded_data.extend(huffman_table[symbol])

    with open(output_file, 'wb') as file:
        encoded_data.tofile(file)

    compression_ratio =   1 - (len(data) )/len(encoded_data)
    print(f'Compression ratio: {compression_ratio:.2f}')
    print(f'Compressed file saved as: {output_file}')


def decompress(input_file, output_file):
    with open(input_file, 'rb') as file:
        encoded_data = bitarray.bitarray()
        encoded_data.fromfile(file)

    frequencies = count_frequencies(encoded_data)
    huffman_tree = build_huffman_tree(frequencies)

    current_node = huffman_tree
    decoded_data = ""
    for bit in encoded_data:
        if bit:
            current_node = current_node.right
        else:
            current_node = current_node.left

        if current_node.left is None and current_node.right is None:
            symbol = str(current_node.symbol)
            decoded_data += symbol
            current_node = huffman_tree

    with open(output_file, 'w', encoding='utf-8') as file:
        file.write(decoded_data)

    print(f'Decompressed file saved as: {output_file}')

# 输入和输出文件名
input_file = 'GCA_014851415.1_ASM1485141v1_genomic.fna'
compressed_file = 'compressed_file.bin'
decompressed_file = 'decompressed_file.txt'

# 压缩文件
compress(input_file, compressed_file)

# 解压缩文件
decompress(compressed_file, decompressed_file)